package com.tous.common.config;

import com.google.protobuf.Message;
import com.googlecode.protobuf.format.JsonFormat;
import org.springframework.http.HttpInputMessage;
import org.springframework.http.HttpOutputMessage;
import org.springframework.http.MediaType;
import org.springframework.http.converter.AbstractHttpMessageConverter;
import org.springframework.http.converter.HttpMessageNotReadableException;
import org.springframework.http.converter.HttpMessageNotWritableException;

import java.io.IOException;
import java.lang.reflect.Method;

/**
 * @author 73578
 * @date 2018/11/27 10:30
 * @description  使用protobuf协议
 */
public class ProtobufHttpMessageConverter extends AbstractHttpMessageConverter<Message> {

    public static final MediaType PROTOBUF = new MediaType("application", "x-protobuf");
    public static final String PROTOBUF_HEADER = "X-Protobuf-Schema";
    public static final String PROTOBUF_MESSAGE_HEADER = "X-Protobuf-Message";

    private final JsonFormat jsonFormat;

    public ProtobufHttpMessageConverter() {
        super(PROTOBUF, MediaType.APPLICATION_JSON);
        this.jsonFormat = new JsonFormat();
    }

    @Override
    protected boolean supports(Class<?> clazz) {
        return Message.class.isAssignableFrom(clazz);
    }

    @Override
    protected Message readInternal(Class<? extends Message> clazz, HttpInputMessage inputMessage)
            throws IOException, HttpMessageNotReadableException {
        MediaType contentType = inputMessage.getHeaders().getContentType();
        Message.Builder builder = getMessageBuilder(clazz);
        
        if (contentType != null && contentType.isCompatibleWith(MediaType.APPLICATION_JSON)) {
            jsonFormat.merge(inputMessage.getBody(), builder);
        } else {
            builder.mergeFrom(inputMessage.getBody());
        }
        return builder.build();
    }

    @Override
    protected void writeInternal(Message message, HttpOutputMessage outputMessage)
            throws IOException, HttpMessageNotWritableException {
        MediaType contentType = outputMessage.getHeaders().getContentType();
        
        if (contentType != null && contentType.isCompatibleWith(MediaType.APPLICATION_JSON)) {
            jsonFormat.print(message, outputMessage.getBody());
        } else {
            message.writeTo(outputMessage.getBody());
            outputMessage.getHeaders().set(PROTOBUF_MESSAGE_HEADER, message.getDescriptorForType().getFullName());
        }
    }

    private Message.Builder getMessageBuilder(Class<? extends Message> clazz) {
        try {
            Method method = clazz.getMethod("newBuilder");
            return (Message.Builder) method.invoke(null);
        } catch (Exception e) {
            throw new HttpMessageNotReadableException("Could not create builder for " + clazz, e);
        }
    }
}    